#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"

#include "hyper/Hyper.h"

#include <iostream>

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir


using namespace mlir;
using namespace vector;
using namespace mlir::arith;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

class HyperMemcpyLowering : public OpRewritePattern<hyper::MemcpyOp> {
public:
  using OpRewritePattern<hyper::MemcpyOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(hyper::MemcpyOp op,
                                PatternRewriter &rewriter) const override {
    StringRef dir = op.getCopyDirection();    
    Value src = op.getSrc();
    Value dst = op.getDst();
    // auto srcOp = src.getDefiningOp();
    // auto dstOp = dst.getDefiningOp();
    // // get dialect name of src and dst ops
    // StringRef srcDialect = srcOp->getDialect()->getNamespace();
    // StringRef dstDialect = dstOp->getDialect()->getNamespace();

    // // only support data movement between host and device
    // assert(((!srcDialect.equals("hyper") && dstDialect.equals("hyper")) || 
    //        (srcDialect.equals("hyper") && !dstDialect.equals("hyper"))) &&
    //         "only support data movement between host and device");

    // TODO: get the device type from src and dst ops without the need of copy direction
    if (dir.equals("cpu2cpu")) {
      rewriter.replaceOpWithNewOp<memref::CopyOp>(op, /*source=*/src, /*target=*/dst);
      // if (dstDialect.equals("hyper")) {
      //   // erase 'hyper.dealloc' op if its memref operand is to be deleted
      //   for (auto user : dstOp->getUsers()) {
      //     if (auto deallocOp = dyn_cast_or_null<hyper::DeallocOp>(user)) {
      //       rewriter.eraseOp(deallocOp);
      //     }
      //   }
      //   // replace 'hyper.alloc' op with its memref operand
      //   rewriter.replaceOp(dstOp, {src});
      //   rewriter.eraseOp(op);
      // } 
      // if (srcDialect.equals("hyper")) {
      //   // erase 'hyper.dealloc' op if its memref operand is to be deleted
      //   for (auto user : srcOp->getUsers()) {
      //     if (auto deallocOp = dyn_cast_or_null<hyper::DeallocOp>(user)) {
      //       rewriter.eraseOp(deallocOp);
      //     }
      //   }
      //   // replace 'hyper.alloc' op with its memref operand
      //   rewriter.replaceOp(srcOp, {dst});
      //   rewriter.eraseOp(op);
      // }
    } else if (dir.equals("cpu2gpu") || dir.equals("gpu2cpu")) {
      rewriter.replaceOpWithNewOp<gpu::MemcpyOp>(op, /*asyncToken=*/std::nullopt, /*asyncDependencies=*/ValueRange{}, /*dst=*/dst, /*src=*/src);
    } else {
      return rewriter.notifyMatchFailure(op, "unknow copy direction for hyper.memcpy");
    }

    return success();
  }
};

} // namespace

void populateLowerHyperMemcpyPatterns(RewritePatternSet &patterns) {
  patterns.add<HyperMemcpyLowering>(patterns.getContext());
}

namespace {
class LowerHyperMemcpyPass : public PassWrapper<LowerHyperMemcpyPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerHyperMemcpyPass)
  LowerHyperMemcpyPass() = default;
  LowerHyperMemcpyPass(const LowerHyperMemcpyPass &) {}

  StringRef getArgument() const final { return "lower-hyper-memcpy"; }
  StringRef getDescription() const final { return "Lower hyper.memcpy op to device-specific ops."; }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<func::FuncDialect,
                mlir::LLVM::LLVMDialect,
                memref::MemRefDialect,
                arith::ArithDialect,
                vector::VectorDialect,
                affine::AffineDialect,
                gpu::GPUDialect,
                scf::SCFDialect
                >();
  }
};
} // end anonymous namespace.

void LowerHyperMemcpyPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<
  func::FuncDialect,
  mlir::LLVM::LLVMDialect,
  memref::MemRefDialect,
  arith::ArithDialect,
  affine::AffineDialect,
  gpu::GPUDialect,
  vector::VectorDialect,
  scf::SCFDialect
  >();
  target.addLegalOp<
  ModuleOp,
  func::FuncOp,
  func::ReturnOp
  >();

  RewritePatternSet patterns(context);
  populateLowerHyperMemcpyPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerHyperMemcpyPass() { PassRegistration<LowerHyperMemcpyPass>(); }
